import argparse
from pathlib import Path
from typing import ClassVar
import polars as pl
from pre_base import PrepareBase, DATA_DIR

# URL= 'https://ftp.ncbi.nlm.nih.gov/pub/clinvar/tab_delimited/variant_summary.txt.gz'


class PrepareClinvar(PrepareBase):
    name: ClassVar[str] = "ClinVar"
    header: ClassVar[str] = '##INFO=<ID=CLINVAR,Number=1,Type=String,Description="ClinVar, variant_id:chrom:start:end:revstat:clnsig:type">'
    builder: str = "hg38"

    def read_input(self) -> pl.DataFrame:
        dataframe = pl.read_csv(
            self.input_file,
            separator="\t",
            columns=["VariationID", "Chromosome", "Start", "Stop", "ReviewStatus", "ClinicalSignificance", "Assembly", "Type"],
            schema_overrides={"VariationID": str, "Chromosome": str},
        )
        dataframe = (
            dataframe.filter(
                pl.col("Assembly").eq("GRCh38" if self.builder == "hg38" else "GRCh37") & pl.col("Type").is_in(["Deletion", "Duplication"])
            )
            .with_columns(
                pl.format("chr{}", pl.col("Chromosome")).alias("Chromosome"), pl.col("Type").str.slice(0, 3).str.to_uppercase().alias("Type")
            )
            .with_columns(
                pl.format(
                    "{}:{}:{}:{}:{}:{}:{}",
                    pl.col("VariationID"),
                    pl.col("Chromosome"),
                    pl.col("Start"),
                    pl.col("Stop"),
                    pl.col("ReviewStatus"),
                    pl.col("ClinicalSignificance"),
                    pl.col("Type"),
                ).alias("INFO")
            )
            .select(["Chromosome", "Start", "Stop", "INFO"])
        ).sort(["Chromosome", "Start", "Stop"])
        return dataframe


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Prepare ClinVar")
    parser.add_argument("--input_file", "-i", default=DATA_DIR / "raw" / "variant_summary.txt.gz", type=Path, help="input file")
    parser.add_argument("--outdir", "-o", default=DATA_DIR / "bed", type=Path, help="output directory")
    parser.add_argument("--builder", "-b", default="hg38", choices=("hg19", "hg38"), help="genome builder")
    parser.set_defaults(func=lambda args: PrepareClinvar.model_validate(vars(args)).run())
    args = parser.parse_args()
    args.func(args)

    # f"{variant_id}:{self.record.chrom}:{start}:{end}:{revstat}:{clnsig}"
